/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * This source file is part of the Cangjie project, licensed under Apache-2.0
 * with Runtime Library Exception.
 *
 * See https://cangjie-lang.cn/pages/LICENSE for license information.
 */

#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <sys/mman.h>
#include <dlfcn.h>
#include <stdbool.h>
#define _GNU_SOURCE
#include <sched.h>
#include <sigchain.h>

// The reason why we use independent signal handler and add_special_signal_handler API is that
// the runtime cannot handle this error correctly.
// It will trigger SIGSEGV in unwind procedure.
// To fix this, we use add_special_signal_handler in OHOS so that our handler will be triggered
// before it reaches runtime's handler.
bool sigill_handler(int signal, siginfo_t *info, void *context) {
    if (info->si_code == ILL_ILLOPN) {
        fprintf(stderr, "SIGILL\n");
        exit(132);
    } else {
        fprintf(stderr, "Unexpect SIGILL with code: %d\n", info->si_code);
        exit(-1);
    }
    return true;
}

__attribute__((constructor)) void init() {
    cpu_set_t  mask;
    CPU_ZERO(&mask);
    CPU_SET(0, &mask);
    sched_setaffinity(0, sizeof(mask), &mask);
	puts("cpu setting is done");

    struct signal_chain_action act = { 0 };
    act.sca_flags = SA_RESTART | SA_SIGINFO;
    act.sca_sigaction = &sigill_handler;
    add_special_signal_handler(SIGILL, &act);
    fprintf(stderr, "add_special_signal_handler is done\n");
}

struct ExploitConfig {
    uint64_t returnAddressAddon;
    const char *victimMangleName;
    const char *exploitMangleName;
};

struct FunctionHandle {
    uint64_t exploitedAddr;
    uint64_t victimAddr;
};

static void findFunctionAddr(struct ExploitConfig* conf, struct FunctionHandle *result)
{
    // --link-options='--export-dynamic is needed for dlsym operation
    uint64_t exploitedAddr = (uint64_t) dlsym(RTLD_DEFAULT, conf->exploitMangleName);
    fprintf(stderr, "[-] exploit function (%s) address: 0x%lx\n", conf->exploitMangleName, exploitedAddr);

    uint64_t victimAddr = (uint64_t) dlsym(RTLD_DEFAULT, conf->victimMangleName);
    fprintf(stderr, "[-] victim function (%s) address: 0x%lx\n", conf->victimMangleName, victimAddr);

    result->exploitedAddr = exploitedAddr;
    result->victimAddr = victimAddr;
}

__attribute__((always_inline)) static uint64_t stripPac(uint64_t ptr)
{
    asm volatile ("xpaci %[preg]" : [preg] "+r" (ptr) :);
    return ptr;
}

static void changeRa(uint64_t *fp, uint64_t newRa, uint64_t addon)
{
    fprintf(stderr, "[-] pc: %p = 0x%lx\n", fp - 1, fp[-1]);
    fprintf(stderr, "[-] fp: %p = 0x%lx\n", fp + 0, fp[0]);
    fprintf(stderr, "[-] ra: %p = 0x%lx\n", fp + 1, fp[1]);
    fprintf(stderr, "[+] start exploit...\n");

    uint64_t raAddr = newRa + addon;
    fprintf(stderr, "[+] change victim's ra: 0x%lx -> 0x%lx\n", fp[1], raAddr);
    fp[1] = raAddr;

    uint64_t *callerFp = (uint64_t *) *fp;
    fprintf(stderr, "[+] change caller's pc: 0x%lx -> 0x%lx\n", callerFp[-1], newRa);
    callerFp[-1] = newRa;
}

void ropAttacker(struct ExploitConfig* conf)
{
    struct FunctionHandle handle = {0};
    uint64_t *currentFp = (uint64_t *) __builtin_frame_address(0);
    findFunctionAddr(conf, &handle);
    uint64_t exploitedAddr = handle.exploitedAddr;
    uint64_t victimAddr = stripPac(handle.victimAddr);

    if (exploitedAddr == 0) {
        fprintf(stderr, "[!] %s address is not found\n", conf->exploitMangleName);
        goto addr_not_found;
    }

    if (victimAddr == 0) {
        fprintf(stderr, "[!] %s address is not found\n", conf->victimMangleName);
        goto addr_not_found;
    }

    uint64_t *callerFp = (uint64_t *) *currentFp;
    fprintf(stderr, "[*] Searching for start pc from caller fp: %p\n", callerFp);
    int i = 1;
    while (callerFp != NULL) {
        if (stripPac(callerFp[-1]) == victimAddr) {
            fprintf(stderr, "[-] found victim's start pc at frame %d: %p = 0x%lx\n", i, callerFp - 1, callerFp[-1]);
            fprintf(stderr, "[*] apply return address offset addon: 0x%lx\n", conf->returnAddressAddon);
            changeRa(callerFp, exploitedAddr, conf->returnAddressAddon);
            return;
        }
        callerFp = (uint64_t *) *callerFp;
        i++;
    }

    fprintf(stderr, "[!] start pc position not found, possibly %s is inlined\n", conf->victimMangleName);
    return;

addr_not_found:
    fprintf(stderr, "[!] ensure you pass option `--link-options='--export-dynamic'` and check function's mangle name\n");
    return;
}
